//
// Created by wyh on 24-4-17.
//
#include <fstream>
#include<vector>
#include <mutex>
#include <opencv2/opencv.hpp>
#include "NvInferRuntimeCommon.h"
#include "NvInferRuntime.h"
#include "BYTETracker.h"

#ifndef YOLODETECTOR_H
#define YOLODETECTOR_H


using namespace nvinfer1;
namespace ORB_SLAM3
{
class Tracking;
class YoloBoundingBox{
private:
    cv::Rect2f rect;
    std::string label;
    //This id was used for multi object tracking, but in experiment I found that MOT cost too much time, so I delete the MOT module
    //Right now, this id only used for draw colors of bbox
    int class_id;;
    float score;
    
public:
    YoloBoundingBox(cv::Rect2f input_rect,int class_id, float score, std::string input_label, cv::Mat mask);
    // YoloBoundingBox(float x1, float y1, float x2, float y2, std::string input_label, float score);
    std::string GetLabel(){return this->label;}
    cv::Rect2f GetRect(){return this->rect;}
    float GetScore(){return this->score;}
    int GetId(){return this->class_id;}
    void SetId(int inputId){this->class_id = inputId;}

    std::vector<cv::Mat> Getmask(){return this->mask;}

    float average, median, stdcov, average_person;
    float width, height;

    // std::vector<cv::Mat> masks;
    cv::Mat mask;

    std::vector<float> mvKeysDynam;
    std::vector<float> mvKeysDynamIntensity;
//    unordered_map<int,float> mvKeysDynam;
};


class YoloDetector {

private:
    std::string cuda_post_process = "c";
    std::string engine_name = "/home/wyh/proj/trt/yolov8_master/build/yolov8s-seg.engine";
    std::string labels_filename = "/home/wyh/proj/trt/yolov8_master/coco.names";
    std::vector<std::string> class_names;
    int model_bboxes;
//    std::string img_dir;
//    std::string sub_type = "";
//    float gd = 0.0f, gw = 0.0f;
//    int max_channels = 0;

    IRuntime *runtime = nullptr;
    ICudaEngine *engine = nullptr;
    IExecutionContext *context = nullptr;
    cudaStream_t stream;

    float *device_buffers[3];
    float *output_buffer_host = nullptr;
    float *output_seg_buffer_host = nullptr;
    float *decode_ptr_host=nullptr;
    float *decode_ptr_device=nullptr;
    std::unordered_map<int, std::string> labels_map;
public:
    YoloDetector();
    ~YoloDetector();

    // ByteTrack tracker
    BYTETracker *ptracker;
   // 需要跟踪的类别，可以根据自己需求调整，筛选自己想要跟踪的对象的种类（以下对应COCO数据集类别索引）
    std::vector<int>  trackClasses {0, 1, 2, 3, 5, 7};  // person, bicycle, car, motorcycle, bus, truck
//    std::vector<int>  trackClasses {0};  // person, bicycle, car, motorcycle, bus, truck
    bool isTrackingClass(int class_id);

    // std::vector<STrack> mvStrackslist;

    void DetectByTensorRT(cv::Mat image, std::vector<YoloBoundingBox> &YoloBoundingBoxlist, std::vector<cv::Rect2f>& DynamicArea,std::vector<STrack>&  Strackslist);

    void SetTracker(Tracking* pTracker);

    bool isNewImgArrived();
    void Run();

    bool CheckFinish();
    void RequestFinish();

    std::mutex mMutexNewYoloDetector;
    std::mutex mMutexGetNewImg;
    std::mutex mMutexFinish;
    bool mbNewImgFlag;
    bool mbFinishRequested;
    Tracking* mpTracker;
    cv::Mat mImg;
    bool mbTensorRT;
    bool mbYOLO;



};

}
#endif //YOLODETECTOR_H
